# 赛道巡线跟踪 - By: zyb - 周一 10月 25 2021
from uart import my_uart
from utils import limit_angle
from utils import LEFT, RIGHT, STRAIGHT


class LineDetect:
    def line_track(self, img, line_th=[(53, 100, -128, 127, -128, 127)], roi=(0, 20, 80, 40), err=1, type='default', angle_limit=20):
        """线性回归处理中线 直线和弯道
        """
        direct = STRAIGHT
        angle_deg = 0

        img.binary(line_th, invert=False)  # 二值化图像
        # self.line_width(img)
        line = img.get_regression(line_th, roi=roi, robust=True)
        img.draw_line(0, roi[1], 80, roi[1], color=127)
        img.draw_line(0, roi[3], 80, roi[3], color=127)
        if line:
            if line.magnitude() > 8:
                if type == 'grass' or 'turn':
                    err = 1
                direct, angle_deg = self.line_to_theta_and_dir(line, err=err)

                if type == 'grass':  # 直道修正
                    angle_err = 40 - (line.x1() + line.x2()) / 2
                    angle_deg = int(angle_deg + 1 * angle_err / 2)
                    limit_angle(angle_deg, angle_limit)  # 角度限幅
                    if angle_deg > 1:
                        if (line.x1() + line.x2()) / 2 > 40:
                            direct = RIGHT
                        elif (line.x1() + line.x2()) / 2 < 40:
                            direct = LEFT
                        #else:
                            #direct = STRAIGHT

                elif type == 'turn':  # 弯道修正
                    if (line.x1() + line.x2()) / 2 < 38:  # 走到弯道内测，需要修正 ,外侧不需要修正
                        direct = 1
                        angle_err = 40 - (line.x1() + line.x2()) / 2
                        angle_deg = int((angle_deg + angle_err / 2))
                        if angle_deg < 0:
                            angle_deg = 0
                        if angle_deg > angle_limit:
                            angle_deg = angle_limit
                else:
                    angle_err = 35 - (line.x1() + line.x2()) / 2
                    angle_deg = int(angle_deg +  angle_err / 2)
                    limit_angle(angle_deg, angle_limit)  # 角度限幅
                    if angle_deg > 1:
                        if (line.x1() + line.x2()) / 2 > 35:
                            direct = RIGHT
                        elif (line.x1() + line.x2()) / 2 < 35:
                            direct = LEFT

                my_uart.set_data(direct, 'direction')
                my_uart.set_data(angle_deg, 'angle')
                img.draw_line(line.line(), color=127)
                # print('direct', direct)
                # print('angle_deg', angle_deg)
                return direct, angle_deg

    def line_track_grass(self, img, line_th=[(60, 100, -128, 127, -128, 127)], roi=(0, 20, 80, 40), err=1):
        """ 线性回归处理中线 草地专用 """
        direct = 0
        angle_deg = 0

        img.binary([(15, 60, -67, -18, 12, 50)], invert=False)  # 二值化图像
        img.erode(1)
        line = img.get_regression(line_th, roi=roi, robust=True)
        img.draw_line(0, roi[1], 80, roi[1], color=127)
        img.draw_line(0, roi[3], 80, roi[3], color=127)

        if line:
            if line.magnitude() > 8:
                direct, angle_deg = self.line_to_theta_and_dir(line, err=err)

                img.draw_line(line.line(), color=127)
                angle_err = 40 - (line.x1() + line.x2()) / 2
                angle_deg = int(angle_deg + 1.15 * angle_err / 2)
                limit_angle(angle_deg, 30)
                # TODO： 补充方向修正，可能方向反了
                if angle_deg > 1:
                    if (line.x1() + line.x2()) / 2 > 40:
                        direct = RIGHT
                    elif (line.x1() + line.x2()) / 2 < 40:
                        direct = LEFT

                my_uart.set_data(direct, 'direction')
                my_uart.set_data(angle_deg, 'angle')
                # print('direct', direct)
                # print('angle_deg', angle_deg)
                return direct, angle_deg

    def line_to_theta_and_dir(self, line, err=20):
        """ 计算中线角度 左1中0右2 """
        angle = line.theta()
        angle_deg = 0
        direct = STRAIGHT
        if line.theta() > 90:
            angle_deg = 180 - angle
            if angle_deg < err:
                direct = STRAIGHT
            else:
                direct = LEFT  # 靠左
        elif line.theta() < 90:
            angle_deg = angle
            if angle_deg < err:
                direct = STRAIGHT
            else:
                direct = RIGHT  # 靠右

        return (direct, angle_deg)


my_line = LineDetect()
